from torch.nn import Module
from torch.nn import Sequential
from torch.nn import Conv2d
from torch.nn import Dropout
from torch.nn import ReLU
from torch.nn import MaxPool2d
from torch.nn import Linear
from cat_dog_torch.settings import *


class AlexNet(Module):
    def __init__(self):
        super(AlexNet, self).__init__()

        self.pool = 2 # 最大池化层
        self.padding = 1 # 矩形边的补充层
        self.dropout = 0.5
        self.kernel_size = 3 # 卷积核

        # 卷积池化
        self.layer1 = Sequential(
            # 时序容器Sequential,参数按顺序传入
            # 2维卷积层，卷积核大小为self.kernel_size，边的补充层数为self.padding
            Conv2d(3, 32, kernel_size=self.kernel_size, padding=self.padding),
            # 对输入数据运用修正线性单元函数
            ReLU(),
            # 最大池化
            MaxPool2d(2))

        # 卷积池化
        self.layer2 = Sequential(
            Conv2d(32, 64, kernel_size=self.kernel_size, padding=self.padding),
            ReLU(),
            MaxPool2d(2))

        # 卷积池化
        self.layer3 = Sequential(
            Conv2d(64, 64, kernel_size=self.kernel_size, padding=self.padding),
            ReLU(),
            MaxPool2d(2))

        # 全连接
        self.fc = Sequential(
            Linear((IMAGE_WIDTH // 8) * (IMAGE_HEIGHT // 8) * 64, 1024),
            Dropout(self.dropout),
            ReLU())

        self.rfc = Sequential(
            Linear(1024, 2)
        )


    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = self.layer3(out)
        out = out.view(out.size(0), -1)
        out = self.fc(out)
        out = self.rfc(out)
        return out
